import os
import sys
import json
import requests
import time
from pathlib import Path

import torch
import cv2
from PIL import Image
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor, AutoModel

# 设置中文字体支持
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
matplotlib.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

class MiniCPMVDeploy:
    def __init__(self, config_path="config.json"):
        """初始化MiniCPM-V部署类"""
        # 加载配置文件
        self.config = self.load_config(config_path)
        
        # 从配置中获取模型信息
        model_config = self.config["model"]
        self.model_name = model_config["model_name"]
        self.local_model_path = model_config["local_model_path"]
        self.cache_dir = model_config["cache_dir"]
        
        # 初始化模型相关变量
        self.model = None
        self.tokenizer = None
        self.processor = None
        self.dtype = torch.float32  # 默认为float32
        
        # 创建输出目录
        output_dir = self.config["paths"]["output_dir"]
        os.makedirs(output_dir, exist_ok=True)
        
        # 设置GPU环境
        self.setup_gpu()
    
    def load_config(self, config_path):
        """加载配置文件"""
        if not os.path.exists(config_path):
            print(f"警告: 配置文件{config_path}不存在，将使用默认配置")
            return self.get_default_config()
        
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            print(f"成功加载配置文件: {config_path}")
            return config
        except Exception as e:
            print(f"加载配置文件失败: {e}，将使用默认配置")
            return self.get_default_config()
            
    def get_default_config(self):
        """获取默认配置"""
        return {
            "model": {
                "model_name": "openbmb/MiniCPM-V-4.5",
                "local_model_path": "F:\\AIModel\\MiniCPM-V-4_5-int4",
                "download_timeout": 3600,
                "cache_dir": "models"
            },
            "gpu": {
                "use_gpu": True,
                "device_id": 0,
                "dtype": "float16",
                "max_memory": None,
                "enable_amp": True
            },
            "inference": {
                "max_new_tokens": 1024,
                "temperature": 0.7,
                "top_p": 0.9,
                "do_sample": True
            },
            "video": {
                "sample_interval_seconds": 2,
                "max_frames": 10
            },
            "paths": {
                "default_image_path": "F:\\AI\\MiniCPM-V\\7140912008894909736_4.jpeg",
                "output_dir": "outputs"
            }
        }
        
    def setup_gpu(self):
        """根据配置设置GPU环境"""
        gpu_config = self.config["gpu"]
        
        # 检查是否使用GPU
        if gpu_config["use_gpu"] and torch.cuda.is_available():
            # 设置设备ID
            device_id = gpu_config["device_id"]
            if device_id >= 0 and device_id < torch.cuda.device_count():
                self.device = f"cuda:{device_id}"
                # 设置可见的GPU设备
                os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
            else:
                self.device = "cuda"
                print(f"警告: 指定的GPU设备ID {device_id} 无效，使用默认GPU")
            
            # 设置最大内存（如果配置了）
            if gpu_config["max_memory"]:
                torch.cuda.set_per_process_memory_fraction(gpu_config["max_memory"])
            
            # 设置数据类型
            if gpu_config["dtype"] == "float16":
                self.dtype = torch.float16
            elif gpu_config["dtype"] == "bfloat16" and torch.cuda.is_bf16_supported():
                self.dtype = torch.bfloat16
            else:
                self.dtype = torch.float32
                print(f"警告: 指定的数据类型 {gpu_config["dtype"]} 不支持，使用float32")
            
            print(f"使用GPU设备: {self.device}")
            print(f"数据类型: {self.dtype}")
            print(f"GPU内存: {torch.cuda.get_device_properties(self.device).total_memory / 1024**3:.2f}GB")
        else:
            self.device = "cpu"
            self.dtype = torch.float32
            print("使用CPU设备")
            if gpu_config["use_gpu"] and not torch.cuda.is_available():
                print("警告: 配置要求使用GPU，但系统中未检测到可用的GPU")
        
    def check_local_model(self):
        """检查本地是否有已下载的模型文件"""
        if self.local_model_path and os.path.exists(self.local_model_path):
            # 检查模型文件夹中是否包含必要的文件
            # 检查基本配置文件
            required_config_files = ["config.json", "tokenizer_config.json", "vocab.json"]
            config_files_found = sum(1 for file in required_config_files if os.path.exists(os.path.join(self.local_model_path, file)))
            
            # 检查模型权重文件（支持分割的safetensors文件）
            model_files = []
            if os.path.exists(os.path.join(self.local_model_path, "model.safetensors.index.json")):
                # 分割的模型文件
                model_files.append(True)
                print("发现分割的模型权重文件")
            elif os.path.exists(os.path.join(self.local_model_path, "model.safetensors")):
                # 单一的模型文件
                model_files.append(True)
                print("发现单一的模型权重文件")
            
            # 检查是否有部分配置文件和模型文件
            if config_files_found >= 2 and model_files:
                print(f"发现完整的本地模型: {self.local_model_path}")
                return True
            else:
                print(f"本地模型文件结构不符合要求")
                print(f"配置文件数量: {config_files_found}/{len(required_config_files)}")
                print(f"模型文件: {'存在' if model_files else '不存在'}")
                return False
        else:
            # 检查默认模型文件夹
            models_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), self.cache_dir)
            default_model_path = os.path.join(models_dir, self.model_name.split("/")[-1])
            if os.path.exists(default_model_path):
                self.local_model_path = default_model_path
                print(f"发现本地模型: {self.local_model_path}")
                return True
        return False
    
    def download_model(self):
        """从魔塔社区下载模型"""
        # 创建models文件夹
        models_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
        os.makedirs(models_dir, exist_ok=True)
        
        # 设置本地模型路径
        model_folder_name = self.model_name.split("/")[-1]
        self.local_model_path = os.path.join(models_dir, model_folder_name)
        
        print(f"开始从魔塔社区下载模型到: {self.local_model_path}")
        
        # 使用huggingface_hub下载模型
        try:
            from huggingface_hub import snapshot_download
            snapshot_download(
                repo_id=self.model_name,
                local_dir=self.local_model_path,
                local_dir_use_symlinks=False,
                token=None  # 如果需要登录，可以设置token
            )
            print("模型下载完成")
            return True
        except Exception as e:
            print(f"模型下载失败: {e}")
            # 如果失败，尝试使用transformers直接下载
            try:
                print("尝试使用transformers直接下载并保存模型...")
                # 直接使用AutoTokenizer和AutoProcessor下载并保存模型
                AutoTokenizer.from_pretrained(
                    self.model_name,
                    trust_remote_code=True,
                    cache_dir=models_dir
                )
                AutoProcessor.from_pretrained(
                    self.model_name,
                    trust_remote_code=True,
                    cache_dir=models_dir
                )
                print("transformers下载完成，模型已保存到缓存目录")
                # 更新local_model_path为huggingface默认缓存路径
                import transformers
                cache_dir = transformers.utils.hub.default_cache_path
                self.local_model_path = os.path.join(cache_dir, "models--" + self.model_name.replace("/", "--"))
                return True
            except Exception as e:
                print(f"transformers下载失败: {e}")
                return False
    
    def load_model(self):
        """加载模型 - 根据官方示例使用AutoModel，完全使用本地模型"""
        if not self.check_local_model():
            print(f"本地模型检查失败: {self.config['model']['local_model_path']}")
            print("请确保模型文件完整，或调整config.json中的local_model_path路径")
            return False
        
        print(f"开始加载模型: {self.config['model']['local_model_path']}")
        print(f"强制使用本地模型文件，不进行在线下载")
        
        try:
            # 设置模型加载参数 - 强制使用本地文件
            model_kwargs = {
                "dtype": self.dtype,
                "trust_remote_code": True,
                "local_files_only": True  # 强制只使用本地文件
            }
            
            # 记录开始时间
            start_time = time.time()
            
            # 根据官方示例，使用AutoModel而不是AutoModelForCausalLM
            print(f"正在从本地路径加载模型: {self.config['model']['local_model_path']}")
            self.model = AutoModel.from_pretrained(
                self.config["model"]["local_model_path"],
                **model_kwargs
            )
            
            # 确保模型加载到正确的设备并设置为eval模式
            print(f"将模型移至设备: {self.device}")
            self.model = self.model.eval().to(self.device)
            
            # 加载分词器 - 仅从本地路径加载
            print(f"正在从本地路径加载分词器: {self.config['model']['local_model_path']}")
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.config["model"]["local_model_path"],
                trust_remote_code=True,
                local_files_only=True  # 强制只使用本地文件
            )
            
            # 根据官方示例，可能不需要processor
            self.processor = None
            
            # 记录模型加载时间
            end_time = time.time()
            print(f"模型加载完成，耗时: {end_time - start_time:.2f}秒")
            print(f"模型设备: {next(self.model.parameters()).device}")
            
            return True
        except Exception as e:
            print(f"模型加载失败: {e}")
            import traceback
            print(f"错误详情: {traceback.format_exc()}")
            
            # 不尝试从Hugging Face下载，仅使用本地模型
            print("按照要求，不尝试从Hugging Face下载，仅使用本地模型")
            print("请检查本地模型路径和文件完整性")
            return False
    
    def infer_image(self, image_path, prompt=None):
        """推理单张图片 - 根据官方示例使用model.chat()方法"""
        if not self.model or not self.tokenizer:
            print("模型未加载，请先调用load_model()")
            return None
        
        try:
            # 检查模型设备
            print(f"模型当前设备: {next(self.model.parameters()).device}")
            
            # 加载图片
            print(f"尝试加载图片: {image_path}")
            image = Image.open(image_path).convert("RGB")
            print(f"成功加载图片，尺寸: {image.size}")
            
            # 预处理图片 - 调整大小以避免处理大图片时的问题
            target_size = 1024  # 设置目标尺寸，官方示例没有明确限制，但我们保持调整
            width, height = image.size
            if width > target_size or height > target_size:
                # 计算缩放比例
                scale = target_size / max(width, height)
                new_width = int(width * scale)
                new_height = int(height * scale)
                print(f"调整图片尺寸: {image.size} -> ({new_width}, {new_height})")
                image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            # 使用默认提示词或传入的提示词
            if prompt is None:
                prompt = "请详细描述这张图片"
            print(f"使用提示词: {prompt}")
            
            # 根据官方示例，使用messages格式
            print(f"使用官方示例的messages格式进行推理...")
            msgs = [{'role': 'user', 'content': [image, prompt]}]
            
            # 获取推理配置
            inference_config = self.config["inference"]
            print(f"使用推理配置: {inference_config}")
            
            # 推理 - 使用官方示例的model.chat()方法
            print(f"开始推理...")
            with torch.no_grad():
                start_time = time.time()
                # 如果启用了自动混合精度
                if self.config["gpu"]["enable_amp"] and self.device.startswith("cuda"):
                    print("使用自动混合精度(AMP)")
                    with torch.autocast(device_type=self.device.split(":")[0], dtype=self.dtype):
                        result = self.model.chat(
                            msgs=msgs,
                            tokenizer=self.tokenizer,
                            enable_thinking=inference_config.get("enable_thinking", False),
                            max_new_tokens=inference_config["max_new_tokens"],
                            temperature=inference_config["temperature"],
                            top_p=inference_config["top_p"]
                        )
                else:
                    print("使用标准精度")
                    result = self.model.chat(
                        msgs=msgs,
                        tokenizer=self.tokenizer,
                        enable_thinking=inference_config.get("enable_thinking", False),
                        max_new_tokens=inference_config["max_new_tokens"],
                        temperature=inference_config["temperature"],
                        top_p=inference_config["top_p"]
                    )
                end_time = time.time()
            
            print(f"推理完成，耗时: {end_time - start_time:.2f}秒")
            print(f"推理结果: {result}")
            
            # 保存结果到文件
            output_file = os.path.join(
                self.config["paths"]["output_dir"],
                f"image_result_{os.path.basename(image_path).split('.')[0]}.txt"
            )
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(f"图片路径: {image_path}\n")
                f.write(f"提示词: {prompt}\n")
                f.write(f"推理时间: {end_time - start_time:.2f}秒\n")
                f.write("\n===== 推理结果 =====\n")
                f.write(result)
            print(f"结果已保存到: {output_file}")
            
            return result
        except Exception as e:
            import traceback
            print(f"图像推理失败: {e}")
            print(f"错误详情: {traceback.format_exc()}")
            return None
    
    def infer_video(self, video_path, prompt=None):
        """推理视频 - 根据官方示例使用model.chat()方法"""
        if not self.model or not self.tokenizer:
            print("模型未加载，请先调用load_model()")
            return None
        
        try:
            # 检查模型设备
            print(f"模型当前设备: {next(self.model.parameters()).device}")
            
            print(f"正在推理视频: {video_path}")
            
            # 打开视频文件
            cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                print(f"无法打开视频文件: {video_path}")
                return None
            
            # 获取视频信息
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = frame_count / fps if fps > 0 else 0
            
            print(f"视频信息: {frame_count}帧, {fps:.2f}FPS, 时长约{duration:.2f}秒")
            
            # 从配置中获取采样参数
            video_config = self.config["video"]
            sample_interval = max(1, int(fps * video_config["sample_interval_seconds"]))
            max_frames = video_config["max_frames"]
            
            # 采样关键帧
            frames = []
            for i in range(frame_count):
                ret, frame = cap.read()
                if not ret:
                    break
                
                # 采样关键帧
                if i % sample_interval == 0 or i == frame_count - 1:
                    # 转换为RGB并添加到列表
                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    frames.append(Image.fromarray(frame_rgb))
                    
                    if len(frames) >= max_frames:
                        break
            
            cap.release()
            print(f"采样了{len(frames)}帧进行分析")
            
            # 使用默认提示词或传入的提示词
            if prompt is None:
                prompt = "请详细描述这个视频的内容"
            print(f"使用提示词: {prompt}")
            
            # 根据官方示例，使用messages格式
            print(f"使用官方示例的messages格式进行视频推理...")
            # 对于视频，我们将关键帧作为消息内容的一部分
            msgs = [{'role': 'user', 'content': [frames[0]] + [prompt]}]
            
            # 获取推理配置
            inference_config = self.config["inference"]
            print(f"使用推理配置: {inference_config}")
            
            # 推理 - 使用官方示例的model.chat()方法
            print(f"开始推理...")
            with torch.no_grad():
                start_time = time.time()
                # 如果启用了自动混合精度
                if self.config["gpu"]["enable_amp"] and self.device.startswith("cuda"):
                    print("使用自动混合精度(AMP)")
                    with torch.autocast(device_type=self.device.split(":")[0], dtype=self.dtype):
                        result = self.model.chat(
                            msgs=msgs,
                            tokenizer=self.tokenizer,
                            enable_thinking=inference_config.get("enable_thinking", False),
                            max_new_tokens=inference_config["max_new_tokens"],
                            temperature=inference_config["temperature"],
                            top_p=inference_config["top_p"]
                        )
                else:
                    print("使用标准精度")
                    result = self.model.chat(
                        msgs=msgs,
                        tokenizer=self.tokenizer,
                        enable_thinking=inference_config.get("enable_thinking", False),
                        max_new_tokens=inference_config["max_new_tokens"],
                        temperature=inference_config["temperature"],
                        top_p=inference_config["top_p"]
                    )
                end_time = time.time()
            
            print(f"视频推理完成，耗时: {end_time - start_time:.2f}秒")
            print(f"推理结果: {result}")
            
            # 保存结果到文件
            output_file = os.path.join(
                self.config["paths"]["output_dir"],
                f"video_result_{os.path.basename(video_path).split('.')[0]}.txt"
            )
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(f"视频路径: {video_path}\n")
                f.write(f"视频信息: {frame_count}帧, {fps:.2f}FPS, 时长约{duration:.2f}秒\n")
                f.write(f"采样帧数: {len(frames)}\n")
                f.write(f"提示词: {prompt}\n")
                f.write(f"推理时间: {end_time - start_time:.2f}秒\n")
                f.write("\n===== 推理结果 =====\n")
                f.write(result)
            print(f"结果已保存到: {output_file}")
            
            return result
        except Exception as e:
            print(f"视频推理失败: {e}")
            return None

def main():
    """主函数"""
    print("===== MiniCPM-V 本地化部署推理工具 ======")
    
    # 初始化部署类，从配置文件加载设置
    deploy = MiniCPMVDeploy()
    
    # 加载模型
    if not deploy.load_model():
        print("模型加载失败，程序退出")
        return
    
    print("\n===== 模型加载成功，开始推理测试 =====")
    
    # 图像推理测试
    # 从配置文件获取测试图片路径
    image_path = deploy.config["paths"]["default_image_path"]
    if os.path.exists(image_path):
        print(f"\n[图像推理测试] 使用配置文件中的图片: {image_path}")
        deploy.infer_image(image_path)
    else:
        print(f"\n[图像推理测试] 未找到配置文件中的图片: {image_path}")
        # 询问用户是否要输入其他图片路径
        print("请输入其他图片路径（留空跳过图像测试）: ")
        custom_image_path = input().strip()
        if custom_image_path and os.path.exists(custom_image_path):
            deploy.infer_image(custom_image_path)
        else:
            print("跳过图像推理测试")
    
    # 视频推理测试 - 使用用户提供的视频路径
    video_path = "F:\\AI\\MiniCPM-V\\main_20250529062545_20250529062745_H265.mp4"
    print(f"\n[视频推理测试] 使用用户提供的视频路径: {video_path}")
    if os.path.exists(video_path):
        deploy.infer_video(video_path)
    else:
        print(f"未找到视频文件: {video_path}")
        print("请输入其他视频路径（留空跳过）: ")
        custom_video_path = input().strip()
        if custom_video_path and os.path.exists(custom_video_path):
            deploy.infer_video(custom_video_path)
        else:
            print("跳过视频推理测试")
    
    print("\n===== 推理测试完成 ======")

if __name__ == "__main__":
    main()